import jax
from rwkv_basic import rwkv_net
from rwkv_batch import rwkv_net_batch
from rwkv_utils import (
    get_tokenizer,
    parse_rwkv_weight,
    rnn_generate,
    rnn_generate_batch_stateless,
)

# choose device(gpu for default)
# jax.config.update('jax_platform_name', 'cpu')

tokenizer = get_tokenizer()
weights_tree = parse_rwkv_weight("weights/RWKV-4-Pile-169M-20220807-8023.pth")
# weights_tree = parse_rwkv_weight(
#     "weights/RWKV-4-World-CHNtuned-3B-v1-20230625-ctx4096.pth"
# )
weights_tree = jax.tree_map(lambda x: jax.device_put(x), weights_tree)

# import time

# start = time.time()
state = None
while 1:
    req = input("Question: ")
    prompt = f"Question: {req}\n\nAnswer:"
    print("\nAnswer:", end='', flush=True)
    # test serial version
    state = rnn_generate(rwkv_net, weights_tree, prompt, n_tokens=200, tokenizer=tokenizer, state=state)

# end = time.time()

# print(f"Total inference time: {end-start:.2f}s")
# test batch version
# rnn_generate_batch_stateless(rwkv_net_batch, weights_tree, prompt, n_tokens=100, tokenizer=tokenizer)
